#ifndef ADVECTION_HPP
#define ADVECTION_HPP

#include "MOL.hpp"

using namespace std;
using namespace Eigen;
double f0(const double &x)
{
    return exp(-20.0 * pow(x - 2, 2)) + exp(-1.0 * pow(x - 5, 2));
}

class advection_equ
{
private:
    double T = 17;
    double h = 0.05;
    double k, mu, a = 1;
    int size, m; // nk(t) ih(x)
    VectorXd U, U0;
    vector<VectorXd> solution;

public:
    advection_equ() : k(0.8 * h)
    {
        mu = a * k / h;
        m = 700;
        size = T / k + 1;

        U = VectorXd::Zero(m + 1);
        for (int i = 0; i <= m; i++)
            U(i) = f0(i * h);
        U0 = U;
        solution.push_back(U);
    }
    void plot()
    {
        ofstream outfile;
        outfile.open("output.m");
        outfile << "x = 0:0.05:" << m * h << ";" << endl;
        outfile << "y = [";
        for (int i = 0; i <= m; i++)
        {
            outfile << U(i) << ",";
        }
        outfile << "];" << endl;
        outfile << "y1 = [";
        for (int i = 0; i <= m; i++)
        {
            outfile << U0(i) << ",";
        }
        outfile << "];" << endl;
        outfile << "plot(x, y,\".\",x, y1, \".\");";

        outfile << "plot(x, y,\".-\");";
        outfile.close();
    }

    VectorXd FTCS()
    {
        MatrixXd A = MatrixXd::Zero(m + 1, m + 1);
        // set A
        for (int i = 0; i <= m - 1; i++)
        {
            A(i + 1, i) = -1;
            A(i, i + 1) = 1;
        }
        A(0, m) = -1;
        A(m, 0) = 1;
        MatrixXd I = MatrixXd::Identity(m + 1, m + 1);
        for (int i = 1; i < size; i++)
        {
            U = (I - mu / 2 * A) * U;
            solution.push_back(U);
        }
        return U;
    }

    VectorXd leapfrog()
    {
        MatrixXd A = MatrixXd::Zero(m + 1, m + 1);
        // set A
        for (int i = 0; i <= m - 1; i++)
        {
            A(i + 1, i) = -1;
            A(i, i + 1) = 1;
        }
        A(0, m) = -1;
        A(m, 0) = 1;
        MatrixXd I = MatrixXd::Identity(m + 1, m + 1);
        VectorXd U1 = U;
        U = (I - mu / 2 * A) * U;
        solution.push_back(U);
        for (int i = 1; i < size; i++)
        {
            VectorXd temp = U;
            U = U1 - mu * A * U;
            U1 = temp;
            solution.push_back(U);
        }

        return U;
    }
    VectorXd Lax_Friedrichs()
    {
        MatrixXd A = MatrixXd::Zero(m + 1, m + 1);
        MatrixXd B = MatrixXd::Zero(m + 1, m + 1);
        // set A
        for (int i = 0; i <= m - 1; i++)
        {
            A(i + 1, i) = -1;
            A(i, i + 1) = 1;
        }
        A(0, m) = -1;
        A(m, 0) = 1;
        // set B
        for (int i = 0; i <= m - 1; i++)
        {
            B(i, i) = -2;
            B(i + 1, i) = 1;
            B(i, i + 1) = 1;
        }
        B(m, m) = -2;
        B(0, m) = 1;
        B(m, 0) = 1;

        MatrixXd I = MatrixXd::Identity(m + 1, m + 1);
        for (int i = 1; i < size; i++)
        {
            U = (I - mu / 2 * A + 0.5 * B) * U;
            solution.push_back(U);
        }
        return U;
    }
    VectorXd Lax_Wendorff()
    {
        MatrixXd A = MatrixXd::Zero(m + 1, m + 1);
        MatrixXd B = MatrixXd::Zero(m + 1, m + 1);
        // set A
        for (int i = 0; i <= m - 1; i++)
        {
            A(i + 1, i) = -1;
            A(i, i + 1) = 1;
        }
        A(0, m) = -1;
        A(m, 0) = 1;
        // set B
        for (int i = 0; i <= m - 1; i++)
        {
            B(i, i) = -2;
            B(i + 1, i) = 1;
            B(i, i + 1) = 1;
        }
        B(m, m) = -2;
        B(0, m) = 1;
        B(m, 0) = 1;

        MatrixXd I = MatrixXd::Identity(m + 1, m + 1);
        for (int i = 1; i < size; i++)
        {
            U = (I - mu / 2 * A + pow(mu, 2) / 2 * B) * U;
            solution.push_back(U);
        }
        return U;
    }
    VectorXd upwind()
    {
        if (a > 0)
        {
            MatrixXd A = MatrixXd::Zero(m + 1, m + 1);
            MatrixXd B = MatrixXd::Zero(m + 1, m + 1);
            // set A
            for (int i = 0; i <= m - 1; i++)
            {
                A(i + 1, i) = -1;
                A(i, i + 1) = 1;
            }
            A(0, m) = -1;
            A(m, 0) = 1;
            // set B
            for (int i = 0; i <= m - 1; i++)
            {
                B(i, i) = -2;
                B(i + 1, i) = 1;
                B(i, i + 1) = 1;
            }
            B(m, m) = -2;
            B(0, m) = 1;
            B(m, 0) = 1;

            MatrixXd I = MatrixXd::Identity(m + 1, m + 1);
            for (int i = 1; i < size; i++)
            {
                U = (I - mu / 2 * A + mu / 2 * B) * U;
                solution.push_back(U);
            }
        }
        return U;
    }
    VectorXd Beam_Warming()
    {
        if (a > 0)
        {
            MatrixXd A = MatrixXd::Zero(m + 1, m + 1);
            MatrixXd B = MatrixXd::Zero(m + 1, m + 1);
            // set A
            for (int i = 0; i <= m - 2; i++)
            {
                A(i, i) = 3;
                A(i + 1, i) = -4;
                A(i + 2, i) = 1;
            }
            A(m, m) = 3;
            A(m - 1, m - 1) = 3;
            A(m, m - 1) = -4;
            A(0, m) = -4;
            A(0, m - 1) = 1;
            A(1, m) = 1;
            // set B
            for (int i = 0; i <= m - 2; i++)
            {
                B(i, i) = 1;
                B(i + 1, i) = -2;
                B(i + 2, i) = 1;
            }
            B(m, m) = 1;
            B(m - 1, m - 1) = 1;
            B(m, m - 1) = -2;
            B(0, m) = -1;
            B(0, m - 1) = 1;
            B(1, m) = 1;

            MatrixXd I = MatrixXd::Identity(m + 1, m + 1);
            for (int i = 1; i < size; i++)
            {
                U = (I - mu / 2 * A + pow(mu, 2) / 2 * B) * U;
                solution.push_back(U);
            }
        }
        return U;
    }
};


#endif